package com.taglab.chic.util;

import java.util.concurrent.*;

/**
 * Memoizer based on Java Concurrency in Practice.
 *
 */
public class Memoizer<A, V> implements Computable<A, V> {
  private final ConcurrentMap<A, FutureTask<V>> cache = new ConcurrentHashMap<A, FutureTask<V>>();
  private final Computable<A, V> computable;

  public Memoizer(Computable<A, V> computable) {
    this.computable = computable;
  }

  public final V compute(final A arg) throws InterruptedException {
    while(true) {
      FutureTask<V> future = cache.get(arg);
      if (future == null) {
        Callable<V> callable = new Callable<V>() {
          public V call() throws InterruptedException {
            return computable.compute(arg);
          }
        };
        FutureTask<V> task = createTask(callable);
        future = cache.putIfAbsent(arg, task);
        if (future == null) {
          // Execute if added
          future = task;
          task.run();
        }
      }
      try {
        return future.get();
      } catch (CancellationException e) {
        cache.remove(arg, future);
      } catch (ExecutionException e) {
        throw launderThrowable(e.getCause());
      }
    }
  }

  protected FutureTask<V> createTask(Callable<V> callable) {
    return new FutureTask<V>(callable);
  }

  private RuntimeException launderThrowable(Throwable throwable) {
    if (throwable instanceof RuntimeException) {
      return (RuntimeException) throwable;
    } else if (throwable instanceof Error) {
      throw (Error) throwable;
    } else {
      throw new IllegalStateException("Not unchecked", throwable);
    }
  }

}
